import torch
import torch.nn as nn
import torchvision.models as models


class VGGEncoder(nn.Module):
    """VGG作为编码器"""
    
    def __init__(self, pretrained=True):
        super(VGGEncoder, self).__init__()
        vgg16 = models.vgg16(pretrained=pretrained).features
        
        self.enc1 = nn.Sequential(*vgg16[:5])  # conv1_1, conv1_2
        self.enc2 = nn.Sequential(*vgg16[5:10]) # conv2_1, conv2_2
        self.enc3 = nn.Sequential(*vgg16[10:17]) # conv3_1, conv3_2, conv3_3
        self.enc4 = nn.Sequential(*vgg16[17:24]) # conv4_1, conv4_2, conv4_3
        self.enc5 = nn.Sequential(*vgg16[24:])  # conv5_1, conv5_2, conv5_3

        # 冻结VGG参数
        if pretrained:
            for param in vgg16.parameters():
                param.requires_grad = False

    def forward(self, x):
        enc1 = self.enc1(x)
        enc2 = self.enc2(enc1)
        enc3 = self.enc3(enc2)
        enc4 = self.enc4(enc3)
        enc5 = self.enc5(enc4)
        return [enc1, enc2, enc3, enc4, enc5]


class VGGSegmentation(nn.Module):
    """基于VGG的分割模型"""
    
    def __init__(self, n_classes=2, pretrained=True):
        super(VGGSegmentation, self).__init__()
        self.encoder = VGGEncoder(pretrained=pretrained)
        
        # 解码器
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2),
            nn.ReLU(inplace=True),
            nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2),
            nn.ReLU(inplace=True),
            nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2),
            nn.ReLU(inplace=True),
            nn.ConvTranspose2d(64, 32, kernel_size=2, stride=2),
            nn.ReLU(inplace=True),
            nn.Conv2d(32, n_classes, kernel_size=1)
        )

    def forward(self, x):
        # 编码器输出的特征图列表
        features = self.encoder(x)
        
        # 使用最深层的特征图进行解码
        out = self.decoder(features[-1]) 
        return out


def get_vgg_segmentation(n_classes=2, pretrained=True):
    """获取VGG分割模型实例"""
    return VGGSegmentation(n_classes=n_classes, pretrained=pretrained)


if __name__ == "__main__":
    # 测试模型
    model = get_vgg_segmentation(n_classes=2, pretrained=False) # 测试时pretrained设为False
    x = torch.randn(1, 3, 256, 256)
    output = model(x)
    print(f"输入形状: {x.shape}")
    print(f"输出形状: {output.shape}")
    print(f"模型参数量: {sum(p.numel() for p in model.parameters() if p.requires_grad)}")

